package edu.buaa.energy;

import com.alibaba.fastjson.JSON;
import edu.buaa.common.BenchmarkWriter;
import edu.buaa.common.transaction.AbstractTransaction;
import edu.buaa.energy.transaction.*;
import edu.buaa.utils.Helper;

import java.io.File;
import java.io.IOException;
import java.util.*;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.BaseStream;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;

import static edu.buaa.common.transaction.AbstractTransaction.TxType.*;

/**
 * Generate an instance of benchmark (a iterator/list of transactions) from given arguments.
 */
public class BenchmarkBuilder {
    private static final String benchmarkFile = Helper.mustEnv("BENCHMARK_FULL_PATH");
    private static final String dataPath = Helper.mustEnv("DATA_FULL_PATH");
    private static final String tpRange = Helper.mustEnv("TP_RANGE").split("_")[1]; //20120101
    private static final int BENCHMARK_SIZE = Integer.parseInt(Helper.mustEnv("BENCHMARK_SIZE"));
    private static final String BENCHMARK_DIST = '['+ Helper.mustEnv("BENCHMARK_DIST") + ']';
    private static final int MAX_CONNECTION_CNT = Integer.parseInt(Helper.envOrDefault("MAX_CONNECTION_CNT", "1"));
    private static final int APPEND_TX_SIZE = Integer.parseInt(Helper.envOrDefault("APPEND_TX_SIZE", "100"));
    private static final boolean IS_VALIDATE = "true".equalsIgnoreCase(System.getenv("NEED_RESULT"));
    private static final String startDay = tpRange.split("-")[0];
    private static final String endDay = tpRange.split("-")[1];

    public static void main(String[] args) throws IOException {
        BenchmarkBuilder builder = new BenchmarkBuilder();
        BenchmarkWriter writer = new BenchmarkWriter(new File(benchmarkFile), !IS_VALIDATE);

        writer.write(builder.readWriteMix(builder.distribution()));
        writer.close();
    }

    private List<Integer> distribution(){
        List<Integer> arr = JSON.parseArray(BENCHMARK_DIST, Integer.class);
        return arr.stream().map(i -> i * BENCHMARK_SIZE / 100).collect(Collectors.toList());
    }

    private Iterator<AbstractTransaction> readWriteMix(List<Integer> arr) throws IOException {
//        System.out.println(arr);
        int tBegin = EnergyDataTxGenerator.parseTime(startDay);
        int tFinish = EnergyDataTxGenerator.parseTime(endDay);
        // append, update, entity_history, snapshot, aggr_max, aggr_dur, entity_temporal_condition, reachable_area
        Stream<AbstractTransaction> txAppend = tx(tx_import_temporal_data, arr.get(0));

        Stream<AbstractTransaction> txs = Stream.of(
                tx(tx_update_temporal_data, arr.get(1)),
                tx(tx_query_entity_history, arr.get(2)),
                tx(tx_query_snapshot, arr.get(3)),
                tx(tx_query_snapshot_aggr_max, arr.get(4)),
                tx(tx_query_snapshot_aggr_duration, arr.get(5)),
                tx(tx_query_road_by_temporal_condition, arr.get(6)),
                tx(tx_query_reachable_area, arr.get(7))
                ).flatMap(tx->tx)
                .map(entityId(1514))
                .map(uniformStartTime(tBegin, tFinish, 3600*5, 3600*50))
                .map(propName(EnergyDataTxGenerator.properties[0]))
                .map(updateArgs(1481, 0, 300))
                .map(uniformDurValIntStart(-1, 1000, 3))
                .map(uniformCondValStart(0, 1000, 100));
        if(IS_VALIDATE){
            return txs.flatMap(appendValidate()).flatMap(updateValidate())
                    .map(uniformDurValIntStart(-1, 1000, 3)).iterator();
        }else {
            List<AbstractTransaction> lst = txs.collect(Collectors.toList());
            Collections.shuffle(lst);
            List<AbstractTransaction> appendList = txAppend.collect(Collectors.toList());
            LinkedList<Integer> queue = orderedRandomList(appendList.size(), lst.size());
            List<AbstractTransaction> result = new LinkedList<>();
            Iterator<AbstractTransaction> appendIter = appendList.iterator();
            for(int i=0;i<lst.size();i++){
                while(!queue.isEmpty() && queue.getFirst()==i && appendIter.hasNext()){
                    queue.pollFirst();
                    result.add(appendIter.next());
                }
                result.add(lst.get(i));
            }
            return result.iterator();
        }
    }

    private LinkedList<Integer> orderedRandomList(int size, int max) {
        Random rand = new Random();
        LinkedList<Integer> result = new LinkedList<>();
        while(result.size()<size) result.add(rand.nextInt(max));
        Collections.sort(result);
        return result;
    }

//    private Iterator<AbstractTransaction> update() throws IOException {
//        int tBegin = EnergyDataTxGenerator.parseTime(startDay);
//        int tFinish = EnergyDataTxGenerator.parseTime(endDay);
//
//        Stream<AbstractTransaction> txs = Stream.of(
//                tx(tx_update_temporal_data, 2)
//        ).flatMap(t->t)
//                .map(uniformStartTime(tBegin, tFinish, 3600*5))
//                .map(updateArgs(1481, 0, 300))
//                .flatMap(updateValidate())
//                .map(uniformDurValIntStart(-1, 1000, 3));
//        return txs.iterator();
//    }
//
//    private Iterator<AbstractTransaction> updateOnly() throws IOException {
//        int tBegin = EnergyDataTxGenerator.parseTime(startDay);
//        int tFinish = EnergyDataTxGenerator.parseTime(endDay);
//
//        Stream<AbstractTransaction> txs = tx(tx_update_temporal_data, 100)
//                .map(uniformStartTime(tBegin, tFinish, 3600*5))
//                .map(updateArgs(1481, 0, 300));
////                .flatMap(updateValidate())
////                .map(uniformDurValIntStart(-1, 1000, 3));
//        return txs.iterator();
//    }
//
//    private Iterator<AbstractTransaction> queries(int benchmarkSize) throws IOException {
//        int tBegin = EnergyDataTxGenerator.parseTime(startDay);
//        int tFinish = EnergyDataTxGenerator.parseTime(endDay);
//
//        Stream<AbstractTransaction> txs = Stream.of(
//                tx(tx_query_snapshot, 200),
//                tx(tx_query_snapshot_aggr_max, 200),
//                tx(tx_query_snapshot_aggr_duration, 200),
//                tx(tx_query_road_by_temporal_condition, 200)
////               tx(tx_update_temporal_data, 200)
//        ).flatMap(t->t)
//                .map(uniformStartTime(tBegin, tFinish, 3600*5, 3600*50))
//                .map(propName(EnergyDataTxGenerator.properties[0]))
//                .map(uniformDurValIntStart(-1, 1000, 3))
//                .map(uniformCondValStart(0, 1000, 100));
////               .map(updateArgs(1481, 0, 300));
//        return txs.iterator();
//    }

    static Stream<AbstractTransaction> tx(AbstractTransaction.TxType type, int cnt) throws IOException {
        System.out.println(type +" "+ cnt);
        if(cnt==0) return Stream.empty();
        else if(type==tx_import_temporal_data) {
            return StreamSupport.stream(importTemporalTx(cnt), false);
        } else{
            return StreamSupport.stream(limit(create(type), cnt), false);
        }
    }

    public static Spliterator<AbstractTransaction> limit(Supplier<AbstractTransaction> supplier, int cnt) {
        return new Spliterators.AbstractSpliterator<AbstractTransaction>(cnt, Spliterator.SIZED) {
            private int i = 0;
            @Override
            public boolean tryAdvance(Consumer<? super AbstractTransaction> action) {
                action.accept(supplier.get());
                return ++i < cnt;
            }
        };
    }

    public static Spliterator<AbstractTransaction> importTemporalTx(int cnt) throws IOException {
        return new Spliterators.AbstractSpliterator<AbstractTransaction>(cnt, Spliterator.SIZED) {
            private int i = 0;
            final Iterator<ImportTemporalDataTx> iter = new EnergyDataTxGenerator(new File(dataPath)).readTemporalStatus(APPEND_TX_SIZE, MAX_CONNECTION_CNT, endDay);
            @Override
            public boolean tryAdvance(Consumer<? super AbstractTransaction> action) {
                if(iter.hasNext()) {
                    action.accept(iter.next());
                    return ++i < cnt;
                }
                return false;
            }
        };
    }

    public static Supplier<AbstractTransaction> create(AbstractTransaction.TxType type){
        return () -> {
            switch (type){
                case tx_query_snapshot: return new SnapshotQueryTx();
                case tx_query_snapshot_aggr_max: return new SnapshotAggrMaxTx();
                case tx_query_snapshot_aggr_duration: return new SnapshotAggrDurationTx();
                case tx_query_road_by_temporal_condition: return new EntityTemporalConditionTx();
                case tx_query_entity_history: return new EntityHistoryTx();
//                case tx_query_reachable_area: return new ReachableAreaQueryTx();
//                case tx_import_temporal_data: return iterator.;
                case tx_update_temporal_data: return new UpdateTemporalDataTx();
                default: throw new UnsupportedOperationException();
            }
        };
    }

    @SafeVarargs
    public static Stream<AbstractTransaction> randMix(Stream<AbstractTransaction> ... streams){
        List<Iterator<AbstractTransaction>> arr = Arrays.stream(streams).map(BaseStream::iterator).collect(Collectors.toList());
        return StreamSupport.stream(new Spliterators.AbstractSpliterator<AbstractTransaction>(Integer.MAX_VALUE, Spliterator.SIZED) {
            private final Random random = new Random();
            @Override
            public boolean tryAdvance(Consumer<? super AbstractTransaction> action) {
                while(!arr.isEmpty()){
                    int i = random.nextInt(arr.size());
                    Iterator<AbstractTransaction> iter = arr.get(i);
                    if(iter.hasNext()){
                        action.accept(iter.next());
                        return true;
                    } else{
                        arr.remove(i);
                    }
                }
                return false;
            }
        }, false);
    }

    public static Function<AbstractTransaction, AbstractTransaction> propName(String name){
        return tx -> {
            switch (tx.getTxType()){
                case tx_query_snapshot:
                    ((SnapshotQueryTx) tx).setPropertyName(name);
                    break;
                case tx_query_snapshot_aggr_max:
                    ((SnapshotAggrMaxTx) tx).setP(name);
                    break;
                case tx_query_snapshot_aggr_duration:
                    ((SnapshotAggrDurationTx) tx).setP(name);
                    break;
                case tx_query_road_by_temporal_condition:
                    ((EntityTemporalConditionTx) tx).setP(name);
                    break;
                case tx_query_entity_history:
//                    assert tx instanceof EntityHistoryTx;
                    ((EntityHistoryTx)tx).setProp(name);
            }
            return tx;
        };
    }

    public static Function<AbstractTransaction, AbstractTransaction> uniformStartTime(int timeBegin, int timeEnd, int length){
        Random random = new Random();
        return tx -> {
            float p = random.nextFloat();
            int time = (int) (timeBegin+ (timeEnd-timeBegin)*p);
            switch (tx.getTxType()){
                case tx_query_snapshot:
                    ((SnapshotQueryTx) tx).setTimestamp(time);
                    break;
                case tx_query_snapshot_aggr_max:
                    ((SnapshotAggrMaxTx) tx).setT0(time);
                    ((SnapshotAggrMaxTx) tx).setT1(time+length);
                    break;
                case tx_query_snapshot_aggr_duration:
                    ((SnapshotAggrDurationTx) tx).setT0(time);
                    ((SnapshotAggrDurationTx) tx).setT1(time+length);
                    break;
                case tx_query_road_by_temporal_condition:
                    ((EntityTemporalConditionTx) tx).setT0(time);
                    ((EntityTemporalConditionTx) tx).setT1(time+length);
                    break;
                case tx_update_temporal_data:
                    ((UpdateTemporalDataTx)tx).setStartTime(time);
                    ((UpdateTemporalDataTx)tx).setEndTime(time+length);
                    break;
                case tx_query_entity_history:
                    assert tx instanceof EntityHistoryTx;
                    ((EntityHistoryTx)tx).setBeginTime(time);
                    ((EntityHistoryTx)tx).setEndTime(time+length);
                    break;
            }
            return tx;
        };
    }

    public static Function<AbstractTransaction, AbstractTransaction> uniformStartTime(int timeBegin, int timeEnd, int len1, int len2){
        Random random = new Random();
        return tx -> {
            float p = random.nextFloat();
            int time = (int) (timeBegin+ (timeEnd-timeBegin)*p);
            int length = random.nextBoolean() ? len1 : len2;
            switch (tx.getTxType()){
                case tx_query_snapshot:
                    ((SnapshotQueryTx) tx).setTimestamp(time);
                    break;
                case tx_query_snapshot_aggr_max:
                    ((SnapshotAggrMaxTx) tx).setT0(time);
                    ((SnapshotAggrMaxTx) tx).setT1(time+length);
                    break;
                case tx_query_snapshot_aggr_duration:
                    ((SnapshotAggrDurationTx) tx).setT0(time);
                    ((SnapshotAggrDurationTx) tx).setT1(time+length);
                    break;
                case tx_query_road_by_temporal_condition:
                    ((EntityTemporalConditionTx) tx).setT0(time);
                    ((EntityTemporalConditionTx) tx).setT1(time+length);
                    break;
                case tx_update_temporal_data:
                    ((UpdateTemporalDataTx)tx).setStartTime(time);
                    ((UpdateTemporalDataTx)tx).setEndTime(time+length);
                    break;
                case tx_query_entity_history:
                    assert tx instanceof EntityHistoryTx;
                    ((EntityHistoryTx)tx).setBeginTime(time);
                    ((EntityHistoryTx)tx).setEndTime(time+length);
                    break;
            }
            return tx;
        };
    }

    public static Function<AbstractTransaction, AbstractTransaction> uniformCondValStart(float valMin, float valMax, int length){
        Random random = new Random();
        return tx -> {
            if(tx.getTxType() == tx_query_road_by_temporal_condition){
                float p = random.nextFloat();
                float min = valMin + p * (valMax-valMin);
                ((EntityTemporalConditionTx) tx).setVMin(min);
                ((EntityTemporalConditionTx) tx).setVMax(min+length);
            }
            return tx;
        };
    }

    public static Function<AbstractTransaction, AbstractTransaction> uniformDurValIntStart(float valMin, float valMax, int cnt){
        return tx -> {
            if(tx.getTxType() == AbstractTransaction.TxType.tx_query_snapshot_aggr_duration){
                float step = (valMax - valMin) / cnt;
                List<Float> intStarts = new ArrayList<>();
                for(float s = valMin; s < valMax; s+= step){
                    intStarts.add(s);
                }
                ((SnapshotAggrDurationTx) tx).setIntervalStarts(intStarts);
            }
            return tx;
        };
    }

    public static Function<AbstractTransaction, Stream<AbstractTransaction>> appendValidate(){
        Random rand = new Random();
        String[] props = EnergyDataTxGenerator.properties;
        return tx -> {
            List<AbstractTransaction> list = new LinkedList<>();
            list.add(tx);
            if(tx.getTxType() == tx_import_temporal_data){
                ImportTemporalDataTx t = (ImportTemporalDataTx) tx;
                int time = t.getData().get(0).getTime();
                list.add(snapshot(props[rand.nextInt(props.length)], time));
                list.add(snapAggrMax(props[rand.nextInt(props.length)], time-3600, time+3600));
                list.add(snapAggrDur(props[rand.nextInt(props.length)], time-3600, time+3600));
            }
            return list.stream();
        };
    }

    public static Function<AbstractTransaction, Stream<AbstractTransaction>> updateValidate(){
        Random rand = new Random();
        String[] props = EnergyDataTxGenerator.properties;
        return tx -> {
            List<AbstractTransaction> list = new LinkedList<>();
            if(tx.getTxType() == tx_update_temporal_data){
                list.add(tx);
                UpdateTemporalDataTx t = (UpdateTemporalDataTx) tx;
                int begin = t.getStartTime(), end = t.getEndTime();
                list.add(snapAggrMax(props[rand.nextInt(props.length)], begin, end));
                list.add(snapAggrDur(props[rand.nextInt(props.length)], begin, end));
                list.add(snapAggrMax(props[rand.nextInt(props.length)], begin-3600, end+3600));
                list.add(snapAggrDur(props[rand.nextInt(props.length)], begin-3600, end+3600));
            }else{
                list.add(tx);
            }
            return list.stream();
        };
    }

    public static Function<AbstractTransaction, AbstractTransaction> entityId(int maxEntityId){
        Random rand = new Random();
        Set<Integer> noExist = new HashSet<>(Arrays.asList(759, 1154 , 1193 , 1225 , 1243, 1245 , 1246 , 1247 , 1248 , 1249, 1252 , 1255 , 1256 , 1263 , 1264, 1276 , 1340 , 1399 , 1501 , 1504));
        return tx -> {
            if(tx.getTxType()==tx_query_entity_history){
                EntityHistoryTx t = (EntityHistoryTx) tx;
                int eid = rand.nextInt(maxEntityId - 1) + 1;
                while(noExist.contains(eid)) eid = rand.nextInt(maxEntityId - 1) + 1;
                t.setEntity(eid);
            }
            return tx;
        };
    }

    public static Function<AbstractTransaction, AbstractTransaction> updateArgs(int maxEntityId, float valMin, float valMax){
        Random rand = new Random();
        Set<Integer> noExist = new HashSet<>(Arrays.asList(759, 1154 , 1193 , 1225 , 1243, 1245 , 1246 , 1247 , 1248 , 1249, 1252 , 1255 , 1256 , 1263 , 1264, 1276 , 1340 , 1399 , 1501 , 1504));
        String[] props = EnergyDataTxGenerator.properties;
        return tx -> {
            if(tx.getTxType() == tx_update_temporal_data){
                UpdateTemporalDataTx t = (UpdateTemporalDataTx) tx;
                int eid = rand.nextInt(maxEntityId - 1) + 1;
                while(noExist.contains(eid)) eid = rand.nextInt(maxEntityId - 1) + 1;
                t.setEntityId(eid);
                List<Float> val = new ArrayList<>();
                for(int i=0; i<props.length; i++) {
                    float p = rand.nextFloat();
                    float min = valMin + p * (valMax-valMin);
                    val.add(min);
                }
                t.setValues(val);
            }
            return tx;
        };
    }

    public static AbstractTransaction snapshot(String propName, int time){
        SnapshotQueryTx tx = new SnapshotQueryTx();
        tx.setTimestamp(time);
        tx.setPropertyName(propName);
        return tx;
    }

    public static AbstractTransaction snapAggrMax(String propName, int timeBegin, int timeEnd){
        SnapshotAggrMaxTx tx = new SnapshotAggrMaxTx();
        tx.setP(propName);
        tx.setT0(timeBegin);
        tx.setT1(timeEnd);
        return tx;
    }

    public static AbstractTransaction snapAggrDur(String propName, int timeBegin, int timeEnd, Float... intStarts){
        SnapshotAggrDurationTx tx = new SnapshotAggrDurationTx();
        tx.setP(propName);
        tx.setT0(timeBegin);
        tx.setT1(timeEnd);
        tx.setIntervalStarts(Arrays.asList(intStarts));
        return tx;
    }

    public static AbstractTransaction entityTimeCond(String propName, int timeBegin, int timeEnd, float valMin, float valMax){
        EntityTemporalConditionTx tx = new EntityTemporalConditionTx();
        tx.setP(propName);
        tx.setT0(timeBegin);
        tx.setT1(timeEnd);
        tx.setVMin(valMin);
        tx.setVMax(valMax);
        return tx;
    }

    public static AbstractTransaction update(int entity, int timeBegin, int timeEnd, List<Float> val){
        UpdateTemporalDataTx tx = new UpdateTemporalDataTx();
        tx.setEntityId(entity);
        tx.setStartTime(timeBegin);
        tx.setEndTime(timeEnd);
        tx.setValues(val);
        return tx;
    }


}
